import z3
import numpy as np
from planner_utils import *

def simple_plan(n_steps, single_constraint):
    # number of interpolating points on trajectory, number of steps
    max_step_length = 0.05 # maximum length of each step
    n_substeps = 1  # ensure that max_step_length/n_substeps << size_of_obs
    
    x = [0]
    y = [0]
    for i in range(n_steps):
        x.append(z3.Real('x'+str(i)))
        y.append(z3.Real('y'+str(i)))
    x.append(0)
    y.append(3)
    dx_plus = []
    dx_minus = []
    dy_plus = []
    dy_minus = []
    
    N = len(x)
    for i in range(N-1):
        dx_plus.append(z3.Real('dx'+str(i)+str(i+1)+'+'))
        dx_minus.append(z3.Real('dx'+str(i)+str(i+1)+'-'))
        dy_plus.append(z3.Real('dy'+str(i)+str(i+1)+'+'))
        dy_minus.append(z3.Real('dy'+str(i)+str(i+1)+'-'))
    
    opt = z3.Optimize()
    for t in np.linspace(0, 1, n_substeps + 1):
        opt.add(free_constraint(x, y, t, single_constraint))
    
    opt.add(step_length_constraint(x, y, max_step_length))
    
    for i in range(N-1):
        opt.add(x[i+1]-x[i] == dx_plus[i] - dx_minus[i])
        opt.add(dx_plus[i]>=0)
        opt.add(dx_minus[i]>=0)
        opt.add(y[i+1]-y[i] == dy_plus[i] - dy_minus[i])
        opt.add(dy_plus[i]>=0)
        opt.add(dy_minus[i]>=0)
    d_all = sum(dx_plus + dx_minus + dy_plus + dy_minus)
    opt.minimize(d_all)
    
    r = opt.check()
    if r == z3.sat:
        model = opt.model()
        for i in range(len(x)):
            print_var(x[i], model)
            print_var(y[i], model)
    else:
        print(r)
    x = [get_var(xi, model) for xi in x]
    y = [get_var(yi, model) for yi in y]
    return x, y
if __name__=="__main__":
    simple_plan(400, free_single_easy)
